import os
import PyPDF2
from pathlib import Path
import pytesseract
from pdf2image import convert_from_path
import io
from PIL import Image

def extract_text_from_pdf(pdf_path):
    """从PDF文件中提取文本内容，包括图片中的文字"""
    try:
        text = ''
        print(f"开始处理PDF文件：{pdf_path}")
        
        # 将PDF转换为高分辨率图片
        print("正在将PDF转换为图片...")
        images = convert_from_path(pdf_path, dpi=600)  # 提高DPI以获取更清晰的图像
        print(f"成功转换 {len(images)} 页图片")
        
        # 使用PyPDF2提取文本
        with open(pdf_path, 'rb') as file:
            pdf_reader = PyPDF2.PdfReader(file)
            
            for i, page in enumerate(pdf_reader.pages):
                print(f"正在处理第 {i+1} 页...")
                
                # 提取常规文本
                try:
                    page_text = page.extract_text()
                    print(f"常规文本提取完成，长度：{len(page_text)}")
                except Exception as e:
                    print(f"常规文本提取失败：{str(e)}")
                    page_text = ''
                
                # 图像预处理和OCR
                try:
                    # 转换为灰度图像
                    img = images[i].convert('L')
                    # 提高对比度
                    from PIL import ImageEnhance
                    enhancer = ImageEnhance.Contrast(img)
                    img = enhancer.enhance(2.0)
                    
                    # 应用自适应阈值处理
                    import cv2
                    import numpy as np
                    img_np = np.array(img)
                    img_np = cv2.adaptiveThreshold(
                        img_np,
                        255,
                        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                        cv2.THRESH_BINARY,
                        11,
                        2
                    )
                    img = Image.fromarray(img_np)
                    
                    # OCR处理
                    print("正在进行OCR处理...")
                    ocr_text = pytesseract.image_to_string(
                        img,
                        lang='chi_sim+chi_tra+eng',  # 同时支持简体中文、繁体中文和英文
                        config='--psm 3 --oem 3 -c preserve_interword_spaces=1'  # 优化中文分词和布局分析
                    )
                    print(f"OCR处理完成，识别文本长度：{len(ocr_text)}")
                except Exception as e:
                    print(f"OCR处理失败：{str(e)}")
                    ocr_text = ''
                
                # 合并文本结果
                final_text = page_text.strip() + '\n' + ocr_text.strip()
                
                if not final_text.strip():
                    print(f"警告：第 {i+1} 页未能提取到任何文本")
                else:
                    print(f"第 {i+1} 页成功提取文本，总长度：{len(final_text)}")
                
                text += final_text + '\n\n'
        
        if not text.strip():
            print("警告：整个PDF文件未能提取到任何文本")
            return None
            
        return text
    except Exception as e:
        print(f"处理PDF文件 {pdf_path} 时出错：{str(e)}")
        import traceback
        print(f"错误详情：\n{traceback.format_exc()}")
        return None

def process_directory(input_dir):
    """处理指定目录下的所有PDF文件"""
    # 确保输入路径存在
    input_path = Path(input_dir)
    if not input_path.exists() or not input_path.is_dir():
        print(f"错误：目录 '{input_dir}' 不存在或不是有效目录")
        return
    
    # 创建输出目录
    output_dir = input_path / 'pdf_texts'
    output_dir.mkdir(exist_ok=True)
    
    # 查找所有PDF文件
    pdf_files = list(input_path.glob('**/*.pdf'))
    if not pdf_files:
        print("未找到任何PDF文件")
        return
    
    print(f"找到 {len(pdf_files)} 个PDF文件")
    
    # 处理每个PDF文件
    for i, pdf_file in enumerate(pdf_files, 1):
        print(f"正在处理 ({i}/{len(pdf_files)}): {pdf_file.name}")
        
        # 提取文本
        text = extract_text_from_pdf(pdf_file)
        if text is None:
            continue
        
        # 创建输出文件
        output_file = output_dir / f"{pdf_file.stem}.txt"
        try:
            with open(output_file, 'w', encoding='utf-8') as f:
                f.write(text)
            print(f"已保存文本到：{output_file}")
        except Exception as e:
            print(f"保存文件 {output_file} 时出错：{str(e)}")

def main():
    print("PDF文本提取工具")
    print("-" * 20)
    
    # 获取用户输入的目录
    while True:
        input_dir = input("请输入包含PDF文件的目录路径（按回车使用当前目录）：").strip()
        if not input_dir:
            input_dir = os.getcwd()
        
        if os.path.isdir(input_dir):
            break
        print(f"错误：'{input_dir}' 不是有效目录，请重新输入")
    
    # 处理目录
    print(f"\n开始处理目录：{input_dir}")
    process_directory(input_dir)
    print("\n处理完成！")

if __name__ == "__main__":
    main()